//
// Created by dave on 18-7-31.
//

#include <atomic>
#include <vector>
#include <glog/logging.h>
#include "thread_pool.h"
#include "blocking_queue.h"

template<typename T>
class pooled_data_saver {
public:
    typedef std::function<void(const std::vector<T>&)> save_func;
    pooled_data_saver(save_func func) : pooled_data_saver(3000, 100, 5, func) { }

    pooled_data_saver(int pool_size, int batch_size, int concurrent_write, save_func func)
            : pool_size_(pool_size), queue_(pool_size), batch_size_(batch_size), concurrent_write_(concurrent_write),
              thread_pool_(concurrent_write), save_(func) { }

    bool save(T data) {
        if (queue_.offer(data)) {
            long counter = ++counter_;
            if (counter % 1000 == 0) {
                clock_t t = clock();
                LOG(INFO) << "已接收：" << counter << ", " << (t - start_) / (CLOCKS_PER_SEC / 1000);
            }
            if (++queue_size_ > batch_size_ && write_thread_.load() < concurrent_write_) {
                doSave();
            }
            return true;
        } else {
//            LOG(INFO) << "入队失败";
            if (doSave() > 0) { // 立即请求写入数据库后再试一次
                return save(data);
            }
            return false;
        }
    }

    void flush() {
        doSave();
    }

private:
    int doSave() {
        int tc = write_thread_.load();
        int qs = queue_size_.load();
        // LOG(ERROR) << "write_thread_："  << tc << "queue_size_: " << qs;
        if (tc < concurrent_write_ && qs > 0) {
            std::vector<T> vec = queue_.drain(batch_size_);
            int size = static_cast<int>(vec.size());
            if (size > 0) {
                write_thread_ += 1;
                queue_size_ -= size;
                thread_pool_.enqueue([this, data = std::move(vec)]() -> void {
                    save_(data);
                    write_thread_ -= 1;
                    if (queue_size_.load() >= batch_size_) {
                        doSave();
                    }
                });
            }
            return size;
        }
        return 0;
    }

private:
    clock_t start_ = clock();
    int pool_size_;
    int batch_size_;
    int concurrent_write_;
    blocking_queue<T> queue_;
    std::atomic_int queue_size_{0};
    std::atomic_int write_thread_{0};
    std::atomic_int counter_{0};
    thread_pool thread_pool_;
    save_func save_;
};